import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
import random
import os
import time



class Chart3D():
    def __init__(self,low=-100,up=100,d=2,iter=[2,40,50],func=None):
        #展示的不同迭代次数图片的参数
        self.a1 = iter[0]
        self.a2 = iter[1]
        self.a3 = iter[2]
        # 绘制三维函数图
        fig = plt.figure()
        # plt.rcParams["font.family"] = ["sans-serif"]
        # plt.rcParams["font.sans-serif"] = ['SimHei']

        x1 = np.arange(low, up, d)
        x2 = np.arange(low, up, d)
        X1, X2 = np.meshgrid(x1, x2)
        nSize = x1.shape[0]
        Z = np.zeros([nSize,nSize])
        for i in range(nSize):
            for j in range(nSize):
                X = [X1[i,j],X2[i,j]]
                X = np.array(X)
                Z[i,j] = func(X)
        self.ax1 = fig.add_subplot(2, 2, 1, projection='3d')
        self.ax2 = fig.add_subplot(2, 2, 2, projection='3d')
        self.ax3 = fig.add_subplot(2, 2, 3, projection='3d')
        self.ax4 = fig.add_subplot(2, 2, 4, projection='3d')
        self.ax1.set_xlabel('X1')
        self.ax1.set_ylabel('X2')
        self.ax1.set_zlabel('Z')
        self.ax1.set_title(f'迭代0次',fontproperties="SimHei")

        self.ax1.plot_surface(X1, X2, Z,cmap='rainbow', zorder=1)
        self.ax2.set_xlabel('X1')
        self.ax2.set_ylabel('X2')
        self.ax2.set_zlabel('Z')
        self.ax2.plot_surface(X1, X2, Z, cmap='rainbow', zorder=1)
        self.ax2.set_title(f'迭代{self.a1}次', fontproperties="SimHei")

        self.ax3.set_xlabel('X1')
        self.ax3.set_ylabel('X2')
        self.ax3.set_zlabel('Z')
        self.ax3.plot_surface(X1, X2, Z, cmap='rainbow', zorder=1)
        self.ax3.set_title(f'迭代{self.a2}次', fontproperties="SimHei")

        self.ax4.set_xlabel('X1')
        self.ax4.set_ylabel('X2')
        self.ax4.set_zlabel('Z')
        self.ax4.plot_surface(X1, X2, Z, cmap='rainbow', zorder=1)
        self.ax4.set_title(f'迭代{self.a3}次', fontproperties="SimHei")
       #self.ax.contour(x,y,z,zdir='z',offset=0)
    def __getitem__(self, key):
        return self.__dict__.get(key, "100")
    def chat3dplot(self,x,fitness,name, xz, color):
        a = []
        b = []
        c = []
        for i in range(len(x)):
            a.append(x[i][0])
            b.append(x[i][1])
            c.append(fitness(x[i]))
        self[name].plot(a, b, c, c=color, marker=xz, zorder=10, linestyle='')


# if __name__ == '__main__':
#     ax = Chart3D()
#     ax.chat3dplot()

